import gym
import sys

sys.path.append('./highway-env/scripts/')

sys.path.append('/home/zy/Desktop/RF_decision/RF_decision_demo/highway-env')
import highway_env

sys.path.append('/home/zy/Desktop/RF_decision/RF_decision_demo/rl-agents')
from rl_agents.agents.common.factory import agent_factory

from tqdm import tnrange
from utils import record_videos, show_videos, capture_intermediate_frames

sys.path.append('/home/zy/Desktop/RF_decision/RF_decision_demo/stable-baselines')
from stable_baselines import HER, SAC


env = gym.make("parking-v0")
model = HER('MlpPolicy', env, SAC, n_sampled_goal=4,
            goal_selection_strategy='future',
            verbose=1, buffer_size=int(1e6),
            learning_rate=1e-3,
            gamma=0.9, batch_size=256,
            policy_kwargs=dict(layers=[256, 256, 256]))
model.learn(int(2e4))

# Visualize a few episodes
from IPython import display as ipythondisplay
from pyvirtualdisplay import Display
from gym.wrappers import Monitor
from pathlib import Path
import base64
from tqdm import tnrange

display = Display(visible=0, size=(1400, 900))
display.start()

def show_video():
    html = []
    for mp4 in Path("video").glob("*.mp4"):
        video_b64 = base64.b64encode(mp4.read_bytes())
        html.append('''<video alt="{}" autoplay 
                      loop controls style="height: 400px;">
                      <source src="data:video/mp4;base64,{}" type="video/mp4" />
                 </video>'''.format(mp4, video_b64.decode('ascii')))
    ipythondisplay.display(ipythondisplay.HTML(data="<br>".join(html)))



# Test the policy
env = gym.make("parking-v0")
env = Monitor(env, './video', force=True, video_callable=lambda episode: True)
for episode in tnrange(3, desc="Test episodes"):
    obs, done = env.reset(), False
    env.unwrapped.automatic_rendering_callback = env.video_recorder.capture_frame
    while not done:
        action, _ = model.predict(obs)
        obs, reward, done, info = env.step(action)
env.close()
show_video()